import { useRef } from "react";
import { Pose } from "@mediapipe/pose";
import * as cam from "@mediapipe/camera_utils";
import * as mediapipePose from "@mediapipe/pose";
import { drawConnectors, drawLandmarks } from "@mediapipe/drawing_utils";
import toast from "react-hot-toast";

export function usePostureDetection(
  webcamRef,
  canvasRef,
  corPostureAudio,
  noPostureAudio
) {
  const postureRef = useRef(null);
  const poseRef = useRef(null);
  const cameraRef = useRef(null);

  let goodPosture = null;
  let loaded = false;
  let badPostureCount = 0;
  let lastToastTime = 0;

  function onResults(results, isRecording) {
    if (!loaded && isRecording) {
      loaded = true;
      console.log("HPE model finished loading.");
    }

    if (!results.poseLandmarks) {
      const now = Date.now();
      if (now - lastToastTime > 2000 && isRecording) {
        toast("No pose detected.");
        noPostureAudio.current.play();
        lastToastTime = now;
      }
      postureRef.current = -1;
      return;
    }

    let landmarks = results.poseLandmarks;
    postureRef.current = null;

    const canvasElement = canvasRef.current;
    const canvasCtx = canvasElement.getContext("2d");

    if (isRecording) {
      canvasCtx.save();
      canvasCtx.clearRect(0, 0, canvasElement.width, canvasElement.height);

      canvasCtx.globalCompositeOperation = "source-over";
      drawConnectors(
        canvasCtx,
        results.poseLandmarks,
        mediapipePose.POSE_CONNECTIONS,
        { color: "#c6fcff", lineWidth: 4 }
      );

      drawLandmarks(canvasCtx, results.poseLandmarks, {
        color: "#c6fcff",
        lineWidth: 1,
        radius: 6,
      });

      canvasCtx.restore();

      if (!goodPosture) {
        goodPosture = landmarks;
        toast.success("Posture calibrated!");
      } else {
        if (isBadPosture(landmarks, goodPosture)) {
          badPostureCount++;
          if (badPostureCount >= 60 && isRecording) {
            toast("Correct your posture!");
            corPostureAudio.current.play();
            badPostureCount = 0;
          }
        } else {
          badPostureCount = 0;
        }
      }
    }
  }

  function isBadPosture(currLandmarks, idealLandmarks) {
    let lookingDown =
      currLandmarks[0].y - idealLandmarks[0].y >
      idealLandmarks[9].y - idealLandmarks[0].y;
    let faceIsClose = idealLandmarks[0].z - currLandmarks[0].z > 0.5;
    return lookingDown || faceIsClose;
  }

  function startPoseDetection(isRecordingCallback) {
    const pose = new Pose({
      locateFile: (file) =>
        `https://cdn.jsdelivr.net/npm/@mediapipe/pose/${file}`,
    });

    pose.setOptions({
      modelComplexity: 1,
      smoothLandmarks: true,
      enableSegmentation: false,
      smoothSegmentation: false,
      minDetectionConfidence: 0.5,
      minTrackingConfidence: 0.5,
    });

    pose.onResults((results) => onResults(results, isRecordingCallback()));

    poseRef.current = pose;

    if (webcamRef.current) {
      const camera = new cam.Camera(webcamRef.current.video, {
        onFrame: async () => {
          if (poseRef.current) {
            await poseRef.current.send({ image: webcamRef.current.video });
          }
        },
      });
      camera.start();
      cameraRef.current = camera;
    }
  }

  function stopPoseDetection() {
    if (poseRef.current) {
      poseRef.current = null;
    }
    const canvasElement = canvasRef.current;
    const canvasCtx = canvasElement.getContext("2d");
    canvasCtx.clearRect(0, 0, canvasElement.width, canvasElement.height);
  }

  function reCalibrate() {
    goodPosture = postureRef.current;
    toast.success("Posture Re-calibrated!");
  }

  return {
    startPoseDetection,
    stopPoseDetection,
    reCalibrate,
  };
}
